import os
import sys
import torch
import time
from diffsynth import ModelManager

# 设置日志文件
log_file = "model_loading.log"
with open(log_file, "w") as f:
    try:
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Script started\n")
        f.flush()

        # 设置模型路径
        model_base_path = 'E:/workspace/llm/text2image/models'
        model_name = 'yiwanji/FLUX_xiao_hong_shu_ji_zhi_zhen_shi_V2'
        model_path = os.path.join(model_base_path, model_name)

        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Model path: {model_path}\n")
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Path exists: {os.path.exists(model_path)}\n")
        f.flush()

        # 检查CUDA
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - CUDA available: {torch.cuda.is_available()}\n")
        if torch.cuda.is_available():
            f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - CUDA device: {torch.cuda.get_device_name(0)}\n")
        f.flush()

        # 初始化模型管理器
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Initializing ModelManager...\n")
        model_manager = ModelManager()
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - ModelManager initialized\n")
        f.flush()

        # 尝试加载模型 - 使用正确的参数名称 file_path_list
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Attempting to load models with file_path_list...\n")
        f.flush()

        # 加载模型
        model_manager.load_models(file_path_list=[model_path])
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Models loaded successfully!\n")
        f.flush()

        # 输出成功消息
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Process completed successfully!\n")
        f.flush()

    except Exception as e:
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - Error occurred: {str(e)}\n")
        import traceback
        traceback.print_exc(file=f)
        f.flush()
        sys.exit(1)

print(f"Process completed. Check {log_file} for details.")